# -*- coding: utf-8 -*-
"""
convert pascal voc label to yolo label
可以在convert_annotation函数中指定W_THRE和H_THRE过滤指定大小以下的小目标

"""

import xml.etree.ElementTree as ET
import sys
import argparse
import os


def convert(size, box):
    """
    convert [xmin, xmax, ymin, ymax] to [center_x, center_y, width, height]
    :param size: image size （width, height）
    :param box: rectangle [xmin, xmax, ymin, ymax]
    :return: center_x, center_y, width, height
    """
    scale_w = 1./size[0]
    scale_h = 1./size[1]
    center_x = (box[0] + box[1])/2.0
    center_y = (box[2] + box[3])/2.0
    width = box[1] - box[0]
    height = box[3] - box[2]
    center_x = center_x * scale_w
    center_y = center_y * scale_h
    width = width * scale_w
    height = height * scale_h
    return (center_x, center_y, width, height)


def convert_annotation(xml_file, classes, yolo_label_dir, W_THRE=1, H_THRE=1):
    """
    convert
    :param xml_file:
    :param classes:
    :param yolo_label_dir:
    :return:
    """
    file_name = xml_file.split('/')[-1]
    file_name = file_name.split('.xml')[0]
    yolo_label_file = open(os.path.join(yolo_label_dir, '{}.txt'.format(file_name)), 'w')
    tree = ET.parse(xml_file)
    root = tree.getroot()
    size = root.find('size')
    image_w = int(size.find('width').text)
    image_h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        # difficult = 0
        object_class = obj.find('name').text
        if object_class not in classes or int(difficult) == 1:
            continue
        class_id = classes.index(object_class)
        xml_box = obj.find('bndbox')
        rect = (float(xml_box.find('xmin').text),
                float(xml_box.find('xmax').text),
                float(xml_box.find('ymin').text),
                float(xml_box.find('ymax').text))
        bbox = convert((image_w, image_h), rect)
        # if width or height equals to 0, pass this object
        if bbox[2] * image_w <= W_THRE or bbox[3] * image_h <= H_THRE:
            continue
        yolo_label_file.write(str(class_id) + " " + " ".join([str(a) for a in bbox]) + '\n')
    yolo_label_file.close()


def voc_to_yolo_label(root_path, yolo_label_dir):
    """
    voc label to yolo
    :param root_path: voc root dir
    :param yolo_label_dir: yolo label dir
    :return: None
    """
    # classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
    #            "car", "cat", "chair", "cow", "diningtable", "dog", "horse",
    #            "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
    classes = ["ship", "tank", "plane"]
    yolo_label_dir = os.path.join(yolo_label_dir, "labels")
    if not os.path.exists(yolo_label_dir):
        os.makedirs(yolo_label_dir)
    xml_dir = os.path.join(root_path, 'Annotations/')
    xml_count = len(os.listdir(xml_dir))
    convert_num = 0
    for file in os.listdir(xml_dir):
        xml_file = os.path.join(xml_dir, file)
        convert_annotation(xml_file, classes, yolo_label_dir)
        convert_num = convert_num + 1
        if convert_num % 100 == 0:
            print('completed: {}, total: {}'.format(convert_num, xml_count))
    print('total converted: {}'.format(convert_num))


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        YOLO_LABEL_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/zongshiban_data/2/ccj/dataset/plane_ship_tank_VOC",
                            help="VOC dataset root directory")
        PARSER.add_argument("--yolo_label_dir", type=str,
                            default="/zongshiban_data/2/ccj/dataset/plane_ship_tank_VOC",
                            help="yolo label file directory")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        YOLO_LABEL_DIR = ARGS.yolo_label_dir
    voc_to_yolo_label(ROOT_DIR, YOLO_LABEL_DIR)
